from flask import Flask, request, redirect, jsonify, abort
from models import Base, URLMap
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import string
import random

app = Flask(__name__)

# 数据库配置
engine = create_engine('sqlite:///database.db')
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)
session = Session()

# 配置短码长度
SHORT_CODE_LENGTH = 6

# 生成短码的字符集
CHARACTERS = string.ascii_letters + string.digits

def generate_short_code(length=SHORT_CODE_LENGTH):
    """生成唯一的短码"""
    while True:
        short_code = ''.join(random.choices(CHARACTERS, k=length))
        if not session.query(URLMap).filter_by(short_code=short_code).first():
            return short_code

@app.route('/shorten', methods=['POST'])
def shorten_url():
    data = request.get_json()
    original_url = data.get('url')
    if not original_url:
        return jsonify({"error": "URL is required"}), 400

    # 检查是否已经存在
    existing = session.query(URLMap).filter_by(original_url=original_url).first()
    if existing:
        short_url = request.host_url + existing.short_code
        return jsonify({"short_url": short_url}), 200

    # 生成新的短码
    short_code = generate_short_code()
    new_url = URLMap(original_url=original_url, short_code=short_code)
    session.add(new_url)
    session.commit()

    short_url = request.host_url + short_code
    return jsonify({"short_url": short_url}), 201

@app.route('/<short_code>')
def redirect_short_url(short_code):
    url = session.query(URLMap).filter_by(short_code=short_code).first()
    if url:
        return redirect(url.original_url, code=302)
    else:
        abort(404)

if __name__ == '__main__':
    app.run(debug=True)
